{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "<a href=\"https://colab.research.google.com/drive/1WmPECHLZcidTU6CnwWyaJMgL1e1AClHu?usp=sharing\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"></a>"
      ],
      "metadata": {
        "id": "akAYRb5OkeOv"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "F1IudHPwkbqR"
      },
      "outputs": [],
      "source": [
        "!pip install -qU google-generativeai"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import google.generativeai as genai\n",
        "import getpass\n",
        "import random"
      ],
      "metadata": {
        "id": "pIUWzRIfk5GC"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Get free-tier Google's Gemini API Key here: https://aistudio.google.com/app/apikey"
      ],
      "metadata": {
        "id": "IA6JdcFok9yO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "API_KEY = getpass.getpass(\"Enter your Google API key: \")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vg_XM60ok7Hn",
        "outputId": "7b5693d5-2f83-480c-c6cf-615ed0b589a9"
      },
      "execution_count": 3,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Enter your Google API key: ··········\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "genai.configure(api_key=API_KEY)"
      ],
      "metadata": {
        "id": "5GeOrJiLlDkS"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class RAPAgent:\n",
        "    def __init__(self, environment):\n",
        "        self.model = genai.GenerativeModel(\"gemini-2.0-flash-exp\")\n",
        "        self.environment = environment\n",
        "\n",
        "    def world_model(self, state, action):\n",
        "        \"\"\"LLM simulates state transition\"\"\"\n",
        "        prompt = f\"\"\"Simulate what happens after this action:\n",
        "\n",
        "        Current State: {state}\n",
        "        Action: {action}\n",
        "\n",
        "        Predict new state (JSON format):\"\"\"\n",
        "\n",
        "        response = self.model.generate_content(prompt).text\n",
        "\n",
        "        # Try to parse as dict, fallback to string\n",
        "        try:\n",
        "            import json\n",
        "            new_state = json.loads(response.strip())\n",
        "        except:\n",
        "            # Simple parsing\n",
        "            new_state = state.copy()\n",
        "            if \"increase\" in action.lower():\n",
        "                for key in new_state:\n",
        "                    if isinstance(new_state[key], (int, float)):\n",
        "                        new_state[key] += 10\n",
        "            elif \"decrease\" in action.lower():\n",
        "                for key in new_state:\n",
        "                    if isinstance(new_state[key], (int, float)):\n",
        "                        new_state[key] = max(0, new_state[key] - 10)\n",
        "\n",
        "        return new_state\n",
        "\n",
        "    def reward_model(self, state, goal):\n",
        "        \"\"\"Evaluate how good a state is\"\"\"\n",
        "        prompt = f\"\"\"Rate this state toward the goal (0-10):\n",
        "\n",
        "        Goal: {goal}\n",
        "        Current State: {state}\n",
        "\n",
        "        Score (just number):\"\"\"\n",
        "\n",
        "        response = self.model.generate_content(prompt).text\n",
        "\n",
        "        try:\n",
        "            score = float(response.strip().split()[0])\n",
        "            return min(max(score / 10, 0), 1)\n",
        "        except:\n",
        "            return 0.5\n",
        "\n",
        "    def generate_actions(self, state):\n",
        "        \"\"\"Generate possible actions from current state\"\"\"\n",
        "        prompt = f\"\"\"Given this state, suggest 3 possible actions:\n",
        "\n",
        "      State: {state}\n",
        "\n",
        "      Actions (one per line):\"\"\"\n",
        "\n",
        "        response = self.model.generate_content(prompt).text\n",
        "        actions = [line.strip() for line in response.split(\"\\n\") if line.strip()][:3]\n",
        "        return actions\n",
        "\n",
        "    def mcts_search(self, state, goal, depth=3):\n",
        "        \"\"\"Monte Carlo Tree Search with LLM world model\"\"\"\n",
        "        best_action = None\n",
        "        best_reward = -1\n",
        "\n",
        "        print(f\"🌲 MCTS Planning (depth={depth})...\\n\")\n",
        "\n",
        "        # Generate possible actions\n",
        "        actions = self.generate_actions(state)\n",
        "\n",
        "        for action in actions:\n",
        "            print(f\"  Testing: {action}\")\n",
        "            total_reward = 0\n",
        "\n",
        "            # Simulate multiple rollouts\n",
        "            for rollout in range(2):\n",
        "                sim_state = state.copy()\n",
        "                rollout_reward = 0\n",
        "\n",
        "                # Simulate future steps\n",
        "                for step in range(depth):\n",
        "                    # World model predicts next state\n",
        "                    sim_state = self.world_model(sim_state, action)\n",
        "\n",
        "                    # Reward model evaluates state\n",
        "                    reward = self.reward_model(sim_state, goal)\n",
        "                    rollout_reward += reward\n",
        "\n",
        "                    # Generate next action for continued simulation\n",
        "                    if step < depth - 1:\n",
        "                        next_actions = self.generate_actions(sim_state)\n",
        "                        action = next_actions[0] if next_actions else action\n",
        "\n",
        "                total_reward += rollout_reward\n",
        "\n",
        "            avg_reward = total_reward / 2\n",
        "            print(f\"    → Avg reward: {avg_reward:.2f}\\n\")\n",
        "\n",
        "            if avg_reward > best_reward:\n",
        "                best_reward = avg_reward\n",
        "                best_action = actions[actions.index(action)]\n",
        "\n",
        "        return best_action, best_reward\n",
        "\n",
        "    def execute_action(self, action):\n",
        "        \"\"\"Execute action in real environment\"\"\"\n",
        "        print(f\"⚡ Executing: {action}\")\n",
        "        new_state = self.environment.execute(action)\n",
        "        print(f\"   New state: {new_state}\\n\")\n",
        "        return new_state\n",
        "\n",
        "    def plan_and_execute(self, goal, max_steps=5):\n",
        "        \"\"\"Main RAP loop: simulate, plan, execute\"\"\"\n",
        "        print(f\"\\n{'='*60}\")\n",
        "        print(f\"🎯 Goal: {goal}\")\n",
        "        print(f\"{'='*60}\\n\")\n",
        "\n",
        "        state = self.environment.get_state()\n",
        "        print(f\"📊 Initial state: {state}\\n\")\n",
        "\n",
        "        for step in range(max_steps):\n",
        "            print(f\"{'─'*60}\")\n",
        "            print(f\"STEP {step + 1}/{max_steps}\")\n",
        "            print(f\"{'─'*60}\\n\")\n",
        "\n",
        "            # Planning: Use MCTS with world model\n",
        "            best_action, expected_reward = self.mcts_search(state, goal)\n",
        "\n",
        "            print(f\"🏆 Selected action: {best_action}\")\n",
        "            print(f\"   Expected reward: {expected_reward:.2f}\\n\")\n",
        "\n",
        "            # Execution: Take real action\n",
        "            state = self.execute_action(best_action)\n",
        "\n",
        "            # Evaluate real outcome\n",
        "            actual_reward = self.reward_model(state, goal)\n",
        "            print(f\"📈 Actual reward: {actual_reward:.2f}\\n\")\n",
        "\n",
        "            # Check if goal reached\n",
        "            if actual_reward > 0.8:\n",
        "                print(f\"✅ Goal achieved!\\n\")\n",
        "                break\n",
        "\n",
        "        print(f\"{'='*60}\")\n",
        "        print(f\"📊 Final state: {state}\")\n",
        "        print(f\"{'='*60}\\n\")\n",
        "\n",
        "        return state"
      ],
      "metadata": {
        "id": "5MEJCVPhlFl8"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Simple environment simulator\n",
        "class Environment:\n",
        "    def __init__(self, initial_state):\n",
        "        self.state = initial_state\n",
        "\n",
        "    def get_state(self):\n",
        "        return self.state\n",
        "\n",
        "    def execute(self, action):\n",
        "        \"\"\"Execute action and return new state\"\"\"\n",
        "        new_state = self.state.copy()\n",
        "\n",
        "        # Parse action and update state\n",
        "        action_lower = action.lower()\n",
        "\n",
        "        if \"collect\" in action_lower or \"gather\" in action_lower:\n",
        "            new_state[\"resources\"] = new_state.get(\"resources\", 0) + 15\n",
        "\n",
        "        elif \"build\" in action_lower or \"construct\" in action_lower:\n",
        "            cost = 20\n",
        "            if new_state.get(\"resources\", 0) >= cost:\n",
        "                new_state[\"structures\"] = new_state.get(\"structures\", 0) + 1\n",
        "                new_state[\"resources\"] -= cost\n",
        "\n",
        "        elif \"research\" in action_lower or \"study\" in action_lower:\n",
        "            new_state[\"knowledge\"] = new_state.get(\"knowledge\", 0) + 10\n",
        "\n",
        "        elif \"train\" in action_lower or \"practice\" in action_lower:\n",
        "            new_state[\"skill\"] = new_state.get(\"skill\", 0) + 5\n",
        "\n",
        "        elif \"explore\" in action_lower:\n",
        "            new_state[\"explored\"] = new_state.get(\"explored\", 0) + 1\n",
        "            new_state[\"resources\"] = new_state.get(\"resources\", 0) + 5\n",
        "\n",
        "        elif \"optimize\" in action_lower or \"improve\" in action_lower:\n",
        "            for key in new_state:\n",
        "                if isinstance(new_state[key], (int, float)) and new_state[key] > 0:\n",
        "                    new_state[key] = int(new_state[key] * 1.1)\n",
        "\n",
        "        self.state = new_state\n",
        "        return new_state"
      ],
      "metadata": {
        "id": "V1Vid_HolPIk"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Example 1: Resource Management\n",
        "print(\"=\"*60)\n",
        "print(\"EXAMPLE 1: Resource Management Task\")\n",
        "print(\"=\"*60)\n",
        "\n",
        "env1 = Environment({\"resources\": 10, \"structures\": 0, \"population\": 5})\n",
        "agent1 = RAPAgent(env1)\n",
        "\n",
        "agent1.plan_and_execute(\"Build 3 structures\", max_steps=4)\n",
        "\n",
        "\n",
        "# Example 2: Skill Development\n",
        "print(\"\\n\" + \"=\"*60)\n",
        "print(\"EXAMPLE 2: Skill Development Path\")\n",
        "print(\"=\"*60)\n",
        "\n",
        "env2 = Environment({\"knowledge\": 0, \"skill\": 0, \"experience\": 0})\n",
        "agent2 = RAPAgent(env2)\n",
        "\n",
        "agent2.plan_and_execute(\"Reach 50 skill points\", max_steps=4)\n",
        "\n",
        "\n",
        "# Example 3: Exploration & Expansion\n",
        "print(\"\\n\" + \"=\"*60)\n",
        "print(\"EXAMPLE 3: Exploration Strategy\")\n",
        "print(\"=\"*60)\n",
        "\n",
        "env3 = Environment({\"explored\": 0, \"resources\": 20, \"territory\": 1})\n",
        "agent3 = RAPAgent(env3)\n",
        "\n",
        "agent3.plan_and_execute(\"Explore 3 new areas and gather resources\", max_steps=4)\n",
        "\n",
        "\n",
        "# Example 4: Multi-objective Task\n",
        "print(\"\\n\" + \"=\"*60)\n",
        "print(\"EXAMPLE 4: Complex Multi-step Task\")\n",
        "print(\"=\"*60)\n",
        "\n",
        "env4 = Environment({\n",
        "    \"resources\": 5,\n",
        "    \"knowledge\": 0,\n",
        "    \"structures\": 0,\n",
        "    \"efficiency\": 1.0\n",
        "})\n",
        "agent4 = RAPAgent(env4)\n",
        "\n",
        "agent4.plan_and_execute(\n",
        "    \"Build 2 structures while maintaining resources above 10\",\n",
        "    max_steps=5\n",
        ")\n",
        "\n",
        "print(\"✅ RAP: Reasoning via Planning Complete!\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 312
        },
        "id": "Kf0PRv-7lS__",
        "outputId": "a37706ab-5605-490e-b47f-f3a18a235a30"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "============================================================\n",
            "EXAMPLE 1: Resource Management Task\n",
            "============================================================\n",
            "\n",
            "============================================================\n",
            "🎯 Goal: Build 3 structures\n",
            "============================================================\n",
            "\n",
            "📊 Initial state: {'resources': 10, 'structures': 0, 'population': 5}\n",
            "\n",
            "────────────────────────────────────────────────────────────\n",
            "STEP 1/4\n",
            "────────────────────────────────────────────────────────────\n",
            "\n",
            "🌲 MCTS Planning (depth=3)...\n",
            "\n",
            "  Testing: 1.  Build Structure (cost: 5 resources)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "ONx8TyH5mL9K"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}